import cv2
import numpy as np
from cv2 import dnn_superres
import time
import os
import os.path as osp

BASEDIR = osp.abspath(osp.dirname(__file__))


def apply_EDSR(img_input_dir, img_out_dir):
    if not os.path.exists(img_out_dir):
        os.makedirs(img_out_dir)
    
    # 创建SR对象...
    sr = dnn_superres.DnnSuperResImpl_create()
    # 读入模型
    path_model = osp.join(BASEDIR, "SR_weights", "EDSR_x4.pb")
    sr.readModel(path_model)
    # 设定算法和放大比例
    sr.setModel("edsr", 4)

    names_list = [name for name in os.listdir(img_input_dir) if name.split('.')[-1] in ['png', 'jpg', 'tif', 'bmp']]
    for name in names_list:
        # 读图
        org_img = cv2.imread(os.path.join(img_input_dir, name))

        print(f"processing {name}...")
        # t0 = time.perf_counter()
        result_img = sr.upsample(org_img)

        # 将原图片线性插值放大
        # scale = 4
        # org_BigPic = cv2.resize(org_img, None, fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR)

        # stack_img = np.hstack((org_BigPic, result_img))
        # 输出
        cv2.imwrite(os.path.join(img_out_dir, name), result_img)
        # print(time.perf_counter() - t0)


if __name__ == '__main__':
    apply_EDSR(img_input_dir='D:/TZB/val_plane_cut/val_planes_cut_clsD', 
               img_out_dir='D:/TZB/val_plane_cut/val_planes_cut_clsD/EDSR/')
    